import base64
import datetime
from io import BytesIO
import numpy as np
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import style

from sklearn.model_selection import cross_validate, cross_val_score, train_test_split
from sklearn import preprocessing, svm
from sklearn.linear_model import LinearRegression,BayesianRidge


def feature_engineering(data):
    data['HL_PCT'] = (data['high'] - data['low']) / data['close'] * 100.0
    data['PCT_change'] = (data['close'] - data['open']) / data['open'] * 100.0
    train_data = data[['close', 'HL_PCT', 'PCT_change', 'volume']]
    forecast_col = 'close'
    data.fillna(value=-99999, inplace=True)
    forecast_out = int(np.math.ceil(0.01 * len(data)))
    train_data['label'] = data[forecast_col].shift(-forecast_out)
    return train_data


def prediction_function(data, train_data, forecast_out, scale):
    X = np.array(train_data.drop(['label'], 1))
    X = preprocessing.scale(X)
    X_lately = X[-forecast_out:]
    X = X[:-forecast_out]
    train_data.dropna(inplace=True)
    y = np.array(train_data['label'])

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1)

    clf = BayesianRidge()
    clf.fit(X_train, y_train)
    accuracy = clf.score(X_test, y_test)

    forecast_set = clf.predict(X_lately)

    data['Forecast'] = data['close']

    last_date = data.iloc[-1].date
    last_unix = last_date.timestamp()
    data.index=[x for x in data.date.values]
    # 预测forecast_out天后的
    one_day = 86400
    next_unix = last_unix + one_day
    # data['Forecast'] = np.nan
    for i in forecast_set:
        next_date = datetime.datetime.fromtimestamp(next_unix)
        next_unix += 86400
        data.loc[next_date] = [np.nan for _ in range(len(data.columns) - 1)] + [i]
    predict_data = data.tail(forecast_out+1)
    return clf, accuracy, predict_data


def output_function(clf, accuracy,sources_data, forecast_data):
    # style.use('ggplot')
    # forecast_data.plot()
    # save_file = BytesIO()
    # plt.savefig(save_file, format='png')
    # save_file_base64 = base64.b64encode(save_file.getvalue()).decode('utf8')
    # save_file.close()
    # plt.close()
    # print(sources_data.tail(1)['label'])
    forecast_data['date'] = forecast_data.index
    forecast_data=forecast_data[['date','Forecast']]
    forecast_data.index=[x for x in range(len(forecast_data))]
    return clf,accuracy, sources_data, forecast_data
